// Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
// SPDX-License-Identifier: MIT

import 'katex/dist/katex.min.css';
import { Check, Copy } from 'lucide-react';
import { useMemo, useState } from 'react';
import ReactMarkdown, {
  type Options as ReactMarkdownOptions,
} from 'react-markdown';
import rehypeKatex from 'rehype-katex';
import remarkGfm from 'remark-gfm';
import remarkMath from 'remark-math';

import { rehypeSplitWordsIntoSpans } from '@/components/deer-flow/rehype';
import { Button } from '@/components/ui/button';
import { cn } from '@/lib/utils';

import Image from './image';
import { Tooltip } from './tooltip';

export function Markdown({
  className,
  children,
  style,
  enableCopy,
  animate = false,
  ...props
}: ReactMarkdownOptions & {
  className?: string;
  enableCopy?: boolean;
  style?: React.CSSProperties;
  animate?: boolean;
}) {
  const rehypePlugins = useMemo(() => {
    if (animate) {
      return [rehypeKatex, rehypeSplitWordsIntoSpans];
    }
    return [rehypeKatex];
  }, [animate]);
  return (
    <div
      className={cn(
        className,
        'prose dark:prose-invert prose-p:my-0 prose-img:mt-0 flex flex-col gap-4',
      )}
      style={style}
    >
      <ReactMarkdown
        remarkPlugins={[remarkGfm, remarkMath]}
        rehypePlugins={rehypePlugins}
        components={{
          a: ({ href, children }) => (
            <a href={href} target="_blank" rel="noopener noreferrer">
              {children}
            </a>
          ),
          img: ({ src, alt }) => (
            <a href={src as string} target="_blank" rel="noopener noreferrer">
              <Image className="rounded" src={src as string} alt={alt ?? ''} />
            </a>
          ),
        }}
        {...props}
      >
        {dropMarkdownQuote(processKatexInMarkdown(children))}
      </ReactMarkdown>
      {enableCopy && typeof children === 'string' && (
        <div className="flex">
          <CopyButton content={children} />
        </div>
      )}
    </div>
  );
}

function CopyButton({ content }: { content: string }) {
  const [copied, setCopied] = useState(false);
  return (
    <Tooltip title="Copy">
      <Button
        variant="outline"
        size="sm"
        className="rounded-full"
        onClick={async () => {
          try {
            await navigator.clipboard.writeText(content);
            setCopied(true);
            setTimeout(() => {
              setCopied(false);
            }, 1000);
          } catch (error) {
            console.error(error);
          }
        }}
      >
        {copied ? <Check className="h-4 w-4" /> : <Copy className="h-4 w-4" />}{' '}
      </Button>
    </Tooltip>
  );
}

function processKatexInMarkdown(markdown?: string | null) {
  if (!markdown) return markdown;

  const markdownWithKatexSyntax = markdown
    .replace(/\\\\\[/g, '$$$$') // Replace '\\[' with '$$'
    .replace(/\\\\\]/g, '$$$$') // Replace '\\]' with '$$'
    .replace(/\\\\\(/g, '$$$$') // Replace '\\(' with '$$'
    .replace(/\\\\\)/g, '$$$$') // Replace '\\)' with '$$'
    .replace(/\\\[/g, '$$$$') // Replace '\[' with '$$'
    .replace(/\\\]/g, '$$$$') // Replace '\]' with '$$'
    .replace(/\\\(/g, '$$$$') // Replace '\(' with '$$'
    .replace(/\\\)/g, '$$$$'); // Replace '\)' with '$$';
  return markdownWithKatexSyntax;
}

function dropMarkdownQuote(markdown?: string | null) {
  if (!markdown) return markdown;
  return markdown
    .replace(/^```markdown\n/gm, '')
    .replace(/^```text\n/gm, '')
    .replace(/^```\n/gm, '')
    .replace(/\n```$/gm, '');
}
